import sys
import os
import subprocess
import shutil

def run_command(command):
    """Run a shell command and handle errors."""
    try:
        subprocess.run(command, check=True, shell=True)
    except subprocess.CalledProcessError as e:
        print(f"命令执行失败: {e}")
        sys.exit(1)

def main():
    # 检查是否提供了文件名作为参数
    if len(sys.argv) < 2:
        print(f"用法: {sys.argv[0]} <download.txt 文件路径>")
        sys.exit(1)

    input_file = sys.argv[1]

    # 确保提供的文件存在
    if not os.path.isfile(input_file):
        print(f"错误: {input_file} 文件不存在。")
        sys.exit(1)

    # 逐行读取文件
    with open(input_file, 'r') as file:
        for line in file:
            accession = line.strip()  # 去掉行尾换行符
            if not accession:
                continue

            print(f"正在下载：{accession}")

            # 使用 datasets 命令下载
            zip_filename = f"{accession}.zip"
            datasets_command = f"datasets download genome accession {accession} --include protein,seq-report --filename {zip_filename}"
            run_command(datasets_command)

            # 解压文件
            unzip_command = f"unzip {zip_filename}"
            run_command(unzip_command)

            # 移动文件
            dataset_dir = f"ncbi_dataset/data/{accession}"
            protein_faa = f"{dataset_dir}/protein.faa"
            assembly_report = f"{dataset_dir}/assembly_data_report.jsonl"

            if os.path.exists(protein_faa):
                shutil.move(protein_faa, f"{accession}.faa")
            if os.path.exists(assembly_report):
                shutil.move(assembly_report, f"{accession}.assembly_data_report.jsonl")

            # 删除临时文件
            os.remove("md5sum.txt")
            os.remove("README.md")
            os.remove(zip_filename)
            shutil.rmtree("ncbi_dataset")

if __name__ == "__main__":
    main()